#图片添加噪声  
#将原始图片路径下的所有图片添加高斯噪声（椒盐噪声）到一个新的文件夹中（图片名字不改）
import cv2
import numpy as np
import os
import random

'''
original_image_path = "F:\\Graduate student\\company\\LayoutGen\\image\\"
add_noise_image_path = "F:\\Graduate student\\company\\LayoutGen\\image_noise\\"

if not os.path.exists(add_noise_image_path):
    os.mkdir(add_noise_image_path)
'''

def add_gaussian_noise(image_in, mean=0, var=0.001):
    '''
    Add Gaussian noise to the image_in
    image_in: input image
    mean:  mean value
    var:  variance
    '''
    temp_image = np.array(image_in/255, dtype=float)
    noise = np.random.normal(mean, var**0.5, temp_image.shape)
    out = temp_image + noise
    if np.all(out < 0):
        low_clip = -1.
    else:
        low_clip = 0.
    out = np.clip(out, low_clip, 1.0)
    out = np.uint8(out*255)
    return out

def add_sp_noise(image,prob):
    '''
    add salt and pepper noise to the image
    prob: the percent of the salt and pepper noise in the image
    '''
    output = np.zeros(image.shape,np.uint8)
    thres = 1 - prob 
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            rdn = random.random()
            if rdn < prob:
                output[i][j] = 0
            elif rdn > thres:
                output[i][j] = 255
            else:
                output[i][j] = image[i,j]
    return output

'''
for files in os.walk(original_image_path):
    for file in files[2]:
        print (file + "-->start!")
        img_name = os.path.splitext(file)[0] + '.jpg'
        fileimgpath = original_image_path + img_name
        im = cv2.imread(fileimgpath)
        print(im.shape)
        noise_img = add_gaussian_noise(im,0,0.001)
        #noise_img = add_sp_noise(im,0.01)
        savefilepath = add_noise_image_path + img_name
        cv2.imwrite(savefilepath,noise_img)
'''